#include "common.h"
#include "memory/memory.h"
#include "memory/cache2.h"

extern uint32_t dram_read(hwaddr_t addr, size_t len);
extern void dram_write(hwaddr_t addr, size_t len, uint32_t data);

typedef struct {
    uint8_t data[L2_BLOCK_SIZE];
    uint32_t tag;
    uint8_t valid;
    uint8_t dirty;
} l2_line_t;

typedef struct {
    l2_line_t ways[L2_WAYS];
} l2_set_t;

static struct {
    l2_set_t sets[L2_SETS];
    uint32_t rng_state;
    uint64_t hits, misses, cycles, writebacks;
} g_l2;

static inline uint32_t l2_get_set_index(hwaddr_t addr) {
    return (addr >> L2_OFFSET_BITS) & L2_SET_MASK;
}

static inline uint32_t l2_get_tag(hwaddr_t addr) {
    return (uint32_t)(addr >> (L2_OFFSET_BITS + L2_INDEX_BITS));
}

static inline uint32_t lcg_next2(uint32_t *s){ *s = *s * 1664525u + 1013904223u; return *s; }

void cache2_init(void) {
    uint32_t s, w;
    for (s = 0; s < L2_SETS; s++) {
        for (w = 0; w < L2_WAYS; w++) {
            g_l2.sets[s].ways[w].valid = 0;
            g_l2.sets[s].ways[w].dirty = 0;
            g_l2.sets[s].ways[w].tag = 0;
        }
    }
    g_l2.rng_state = 13579u;
    g_l2.hits = g_l2.misses = g_l2.cycles = g_l2.writebacks = 0;
}

void cache2_stats_reset(void) { g_l2.hits = g_l2.misses = g_l2.cycles = g_l2.writebacks = 0; }

void cache2_stats_dump(void) {
    uint64_t total = g_l2.hits + g_l2.misses;
    double hr = (total ? (double)g_l2.hits / (double)total : 0.0) * 100.0;
    printf("[L2] hits=%llu misses=%llu total=%llu hit_rate=%.2f%% cycles=%llu writebacks=%llu\n",
        (unsigned long long)g_l2.hits,
        (unsigned long long)g_l2.misses,
        (unsigned long long)total,
        hr,
        (unsigned long long)g_l2.cycles,
        (unsigned long long)g_l2.writebacks);
}

static int l2_find_way(uint32_t set_idx, uint32_t tag){
    uint32_t w; l2_set_t *set = &g_l2.sets[set_idx];
    for (w = 0; w < L2_WAYS; w++) {
        l2_line_t *ln = &set->ways[w];
        if (ln->valid && ln->tag == tag) return (int)w;
    }
    return -1;
}

static void l2_writeback(uint32_t set_idx, int way, hwaddr_t aligned_addr){
    l2_line_t *ln = &g_l2.sets[set_idx].ways[way];
    if (ln->valid && ln->dirty) {
        uint32_t wtag = ln->tag;
        hwaddr_t addr = ((hwaddr_t)wtag << (L2_OFFSET_BITS + L2_INDEX_BITS)) | ((hwaddr_t)set_idx << L2_OFFSET_BITS);
        /* write back to DRAM in 16 dwords */
        uint32_t off;
        for (off = 0; off < L2_BLOCK_SIZE; off += 4) {
            uint32_t v = (uint32_t)ln->data[off] | ((uint32_t)ln->data[off+1] << 8) | ((uint32_t)ln->data[off+2] << 16) | ((uint32_t)ln->data[off+3] << 24);
            dram_write(addr + off, 4, v);
        }
        g_l2.writebacks++;
    }
}

static int l2_choose_victim(uint32_t set_idx){
    l2_set_t *set = &g_l2.sets[set_idx];
    uint32_t w;
    for (w = 0; w < L2_WAYS; w++) if (!set->ways[w].valid) return (int)w;
    return (int)(lcg_next2(&g_l2.rng_state) % L2_WAYS);
}

static void l2_fill_line(uint32_t set_idx, int way, hwaddr_t aligned_addr){
    l2_line_t *ln = &g_l2.sets[set_idx].ways[way];
    uint32_t off;
    for (off = 0; off < L2_BLOCK_SIZE; off += 4) {
        uint32_t val = dram_read(aligned_addr + off, 4);
        ln->data[off+0] = (uint8_t)(val & 0xFF);
        ln->data[off+1] = (uint8_t)((val >> 8) & 0xFF);
        ln->data[off+2] = (uint8_t)((val >> 16) & 0xFF);
        ln->data[off+3] = (uint8_t)((val >> 24) & 0xFF);
    }
    ln->tag = l2_get_tag(aligned_addr);
    ln->valid = 1; ln->dirty = 0;
}

/* Exposed to L1: read a whole aligned block into buf */
void cache2_read_block(hwaddr_t aligned_addr, uint8_t *buf){
    uint32_t set_idx = l2_get_set_index(aligned_addr);
    uint32_t tag = l2_get_tag(aligned_addr);
    int way = l2_find_way(set_idx, tag);
    if (way < 0) {
        g_l2.misses++; g_l2.cycles += 200;
        int vic = l2_choose_victim(set_idx);
        l2_writeback(set_idx, vic, aligned_addr);
        l2_fill_line(set_idx, vic, aligned_addr);
        way = vic;
    } else { g_l2.hits++; g_l2.cycles += 2; }
    /* copy 64B to caller buffer */
    memcpy(buf, g_l2.sets[set_idx].ways[way].data, L2_BLOCK_SIZE);
}

/* Exposed to L1: write bytes into L2 (write-allocate, write-back) */
void cache2_write_bytes(hwaddr_t addr, size_t len, uint32_t data){
    uint32_t set_idx = l2_get_set_index(addr);
    uint32_t tag = l2_get_tag(addr);
    uint32_t off = addr & L2_BLOCK_MASK;
    int way = l2_find_way(set_idx, tag);
    if (way < 0) {
        /* write-allocate: bring block from DRAM */
        hwaddr_t aligned = addr & ~((hwaddr_t)L2_BLOCK_MASK);
        int vic = l2_choose_victim(set_idx);
        l2_writeback(set_idx, vic, aligned);
        l2_fill_line(set_idx, vic, aligned);
        way = vic; g_l2.misses++; g_l2.cycles += 200;
    } else { g_l2.hits++; g_l2.cycles += 2; }

    l2_line_t *ln = &g_l2.sets[set_idx].ways[way];
    /* in-place write to line */
    if (len == 1) ln->data[off] = (uint8_t)(data & 0xFF);
    else if (len == 2) { ln->data[off] = (uint8_t)(data & 0xFF); ln->data[off+1] = (uint8_t)((data>>8)&0xFF);} 
    else { ln->data[off] = (uint8_t)(data & 0xFF); ln->data[off+1] = (uint8_t)((data>>8)&0xFF); ln->data[off+2] = (uint8_t)((data>>16)&0xFF); ln->data[off+3] = (uint8_t)((data>>24)&0xFF);} 
    ln->dirty = 1;
}


